import mitt, { type Emitter, type EventType } from 'mitt';

// 定义扩展的 Emitter 类型
interface ExtendedEmitter<Events extends Record<EventType, unknown>> extends Emitter<Events> {
  /**
   * 监听事件指定次数后自动移除
   * @param type 事件类型
   * @param handler 事件处理函数
   * @param count 监听次数，默认为1
   * @returns 用于手动取消监听的函数
   */
  onLimit<Key extends keyof Events>(
    type: Key,
    handler: (event: Events[Key]) => void,
    count?: number
  ): () => void;

  /**
   * 简化的once方法，基于onLimit实现
   * @param type 事件类型
   * @param handler 事件处理函数
   * @returns 用于手动取消监听的函数
   */
  once<Key extends keyof Events>(type: Key, handler: (event: Events[Key]) => void): () => void;
}

// 创建带扩展方法的emitter
const createEmitter = <
  Events extends Record<EventType, unknown> = Record<EventType, unknown>,
>(): ExtendedEmitter<Events> => {
  const emitter = mitt<Events>();

  // 类型断言为扩展的 emitter
  const extendedEmitter = emitter as ExtendedEmitter<Events>;

  /**
   * 监听事件指定次数后自动移除
   */
  extendedEmitter.onLimit = function <Key extends keyof Events>(
    type: Key,
    handler: (event: Events[Key]) => void,
    count: number = 1
  ): () => void {
    // 确保次数为正整数
    const maxCount = Math.max(1, Math.floor(Number(count)));
    let currentCount = 0;

    const wrapper = (event: Events[Key]) => {
      currentCount++;
      // 执行原始处理函数
      handler(event);

      // 达到指定次数则移除监听
      if (currentCount >= maxCount) {
        emitter.off(type, wrapper);
      }
    };

    emitter.on(type, wrapper);

    // 返回取消监听的函数
    return () => emitter.off(type, wrapper);
  };

  /**
   * 简化的once方法，基于onLimit实现
   */
  extendedEmitter.once = function <Key extends keyof Events>(
    type: Key,
    handler: (event: Events[Key]) => void
  ): () => void {
    return this.onLimit(type, handler, 1);
  };

  return extendedEmitter;
};

// 使用示例
export const emitter = createEmitter();
